#pragma once

#include <vector>
#include <memory>
#include <random>
#include <algorithm>
#include <cassert>

template<typename T>
class RandomSampling
{
public:
    RandomSampling(const std::vector<T> &features, const std::vector<T> &labels)
    {
        m_features = features;
        m_labels = labels;

        // 初始化索引向量
        m_indices.resize(m_features.size());
        for (int64_t i = 0; i < m_features.size(); ++i)
        {
            m_indices[i] = i;
        }

        // 使用随机设备生成种子
        std::random_device rd;
        m_rng.seed(rd());

        reset();
    }

    void reset()
    {
        // 打乱索引顺序
        std::shuffle(m_indices.begin(), m_indices.end(), m_rng);
        
        // 重置当前索引位置
        m_current_idx = 0;
    }

    void sampling(int64_t batch_size, std::vector<T*> &features, std::vector<T*> &labels)
    {
        assert(batch_size < m_features.size());

        // 2. 检查是否需要重新洗牌（到达数据末尾时）
        if (m_current_idx + batch_size > m_features.size())
        {
            reset();
        }

        features.resize(batch_size);
        labels.resize(batch_size);
        int64_t end = m_current_idx + batch_size;
        int64_t features_index{};
        for(int64_t index = m_current_idx; index < end; index++)
        {
            features[features_index] = (&(m_features[m_indices[index]]));
            labels[features_index] = (&(m_labels[m_indices[index]]));
            features_index++;
        }
        m_current_idx = end;
    }
private:
    std::vector<T> m_features;      // 数据集
    std::vector<T> m_labels;        // 标签集
    std::vector<int64_t> m_indices; // 索引向量
    std::mt19937 m_rng;              // 随机数生成器
    int64_t m_current_idx;           // 当前索引位置
};
